package com.cee.config;

import com.cee.auditlog.handler.TransactionHandler;
import com.cee.auditlog.handler.TransactionHandlerContext;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.support.DefaultTransactionStatus;

import javax.sql.DataSource;

public class CustomizedDataSourceTransactionManager extends DataSourceTransactionManager {

    private boolean auditTrailRequired;

    public CustomizedDataSourceTransactionManager(DataSource dataSource, boolean auditTrailRequired) {
        super(dataSource);
        this.auditTrailRequired = auditTrailRequired;
    }

    @Override
    protected void doBegin(Object transaction, TransactionDefinition definition) {
        super.doBegin(transaction, definition);

        if (auditTrailRequired) {
            TransactionHandlerContext.setHandler(new TransactionHandler());
        }
    }

    @Override
    protected void doCommit(DefaultTransactionStatus status) {
        super.doCommit(status);

        if (auditTrailRequired) {
            TransactionHandlerContext.getHandler().handleAll();
        }
    }

    @Override
    protected void doRollback(DefaultTransactionStatus status) {
        super.doRollback(status);

        if (auditTrailRequired) {
            TransactionHandlerContext.removeHandler();
        }
    }
}
